'''Utility for connecting Pose Corrective BlendShapes.

Created by:	Patrick Pautler, Arc Productions

Created:	Jan. 2015

----------------------------
Instructions for use:

The purpose of this script is to create a network of nodes to drive a
specific corrective blendShape channel.

It is designed to be run from the script editor in Maya.
The main function that the user will run is "create_hookup_network".

import pose_blend
reload (pose_blend)
pose_blend.create_hookup_network(pose_name, poser_name_list=[], side='lf', angle=90, driver_attr=None, sdk_val=[], blend_node=None, blend_channel=None)

IMPORTANT:
Before running this script, the blendShape node must exist on the relevant geo.
After running the script, the set driven key created will have to be modified by hand.

----------------------------
Sample Usage:

side = 'lf'
blend_node = 'hi_jacket_blendShape'
blend_channel = '%s_jacket_shoulder_fwd_90' % side
pose_name = shoulder_poserFront
pose_blend.create_hookup_network(pose_name, side=side, blend_node=blend_node, blend_channel=blend_channel)

----------------------------
Parameters:

pose_name - the name of the pose. This is used for determining the poser nodes in the
    rig as well as naming new nodes created as part of the network.
    Below is a standard list of pose names for a human biped character.
    Additional poses can be added if necessary.
        
        elbow_140
        shoulder_poserFront           
        shoulder_poserBack
        shoulder_poserUp
        shoulder_poserDown
        clav_40
        wrist_poserfront
        wrist_poserback            
        wrist_poserUp
        wrist_poserDown
        index_dn_90
        mid_dn_90
        ring_dn_90
        pinky_dn_90
        hip_fwd_90
        knee_140

poser_name_list - this allows the user to specify the poser nodes used if necessary.
    Typically the standard or default poser nodes are used.
    Poser nodes exist in the rig and are meant to drive corrective blendShapes.
    Examples: lf_shoulder_poser, rt_wrist_poser

side - 'lf' for left, 'rt' for right.

angle - the angle of the joint when the blendShape channel is fully turned on.
    For wrist and shoulders, this value is used in one the multiplyDivide nodes in the network.
    If a value is not specified, a default value of 90 is used.
    If the value is negative, use a positive value instead.

driver_attr - the attribute on the poser node that drives the set driven key.
    For example: 'rotateZ' or 'rz'.
    Shoulder and Wrist poses are driven with a more complex network, and this
    parameter is not relevant.
    If the user does not specify this parameter, the script will attempt to use
    a standard one. 

sdk_val - a list of the driver values for the Set Driven Keys.
    For example: [0, -90], [0, -70, -140]
    If the user does not specify this parameter, the script will attempt to use
    standard values.    
    
blend_node = the name of the blendShape node.
    If the user does not specify this parameter, the script will create all nodes
    except for the final set_driven key.    

blend_channel = the blend_channel that should be driven by the network.
    If the user does not specify this parameter, the script will create all nodes
    except for the final set_driven key.

----------------------------

'''


#=======================================================================
# Import 
#=======================================================================
import maya.cmds as cmds
import pymel.core as pmc
import sys
import os
import maya.mel


#=======================================================================
# Main Functions
#=======================================================================

def get_standard_pose_list():
    '''
    Get the standard list of poses used for a biped human rig.
    Wrist and Shoulder poses have the word "poser" in the pose name.
    The pose name is used mainly for determining the required poser nodes.
    '''
    
    pose_list = [
        'elbow_140',
        'shoulder_poserFront',           
        'shoulder_poserBack',
        'shoulder_poserUp',
        'shoulder_poserDown',
        'clav_40',
        'wrist_poserfront',
        'wrist_poserback',            
        'wrist_poserUp',
        'wrist_poserDown',
        'index_dn_90',
        'mid_dn_90',
        'ring_dn_90',
        'pinky_dn_90',
        'hip_fwd_90',
        'knee_140', 
        ]
    
    return pose_list
   
def get_posers_for_pose(pose_name):
    '''
    Get the poser nodes needed for a specific pose.
    Examples of pose names:
        shoulder_poserUp
        wrist_poserUp
    Examples of posers:
        clavicle_poser
        hand_loc1
    '''
    
    pose_type = ''
    if 'shoulder' in pose_name:
        pose_type = 'shoulder'
    elif 'wrist' in pose_name:
        pose_type = 'wrist'
    elif 'elbow' in pose_name:
        pose_type = 'elbow'
    elif 'clav' in pose_name:
        pose_type = 'clavicle' 
    elif 'index' in pose_name:
        pose_type = 'index'
    elif 'mid' in pose_name:
        pose_type = 'middle'        
    elif 'ring' in pose_name:
        pose_type = 'ring'        
    elif 'pinkey' in pose_name or 'pinky' in pose_name:
        pose_type = 'pinky'        
    elif 'hip' in pose_name:
        pose_type = 'hip'
    elif 'knee' in pose_name:
        pose_type = 'knee'      
    else:
        pose_type = 'unknown'

    print 'pose_type = ', pose_type
    
    posers_dict = {
        'shoulder' : ['clavicle_poser', 'shoulder_poser', pose_name],
        'wrist' : ['hand_loc1', 'wrist_poser', pose_name],
        'elbow' : ['elbow_poser'],
        'clavicle' : ['clavicle_poser'],  
        'index' : ['index_poser'],
        'middle' : ['middle_poser'],        
        'ring' : ['ring_poser'],
        'pinky' : ['pinkey_poser'],
        'hip' : ['hip_poser'],            
        'knee' : ['knee_poser'],
        'unknown' : [],
        }
        
        
    poser_list = posers_dict[pose_type]
    
    return poser_list

def get_pose_data(pose_name):
    '''
    Get the standard pose data.
    This includes default Set Driven Key Values for the pose.
    In most cases the SDK values will have to be modified by hand afterwards.
    '''
    
    pose_list = get_standard_pose_list()

    # Get the attr that drives the SDK.
    # Note: shoulder and wrist poses use a more complicated setup.
    driver_attr_dict = {
        'elbow_140' : 'ry',
        'shoulder_poserFront' : '',           
        'shoulder_poserBack' : '',
        'shoulder_poserUp' : '',
        'shoulder_poserDown' : '',
        'clav_40' : 'ty',
        'wrist_poserfront' : '',
        'wrist_poserback' : '',            
        'wrist_poserUp' : '',
        'wrist_poserDown' : '',
        'index_dn_90' : 'rz',
        'mid_dn_90' : 'rz',
        'ring_dn_90' : 'rz',
        'pinky_dn_90' : 'rz',
        'hip_fwd_90' : 'tz',
        'knee_140' : 'ry',
        }
    
    # Get the default SDK values.
    # These will probably need to be modified by hand.
    sdk_dict = {
        'elbow_140' : [0.0, -70.0, -140.0],
        'shoulder_poserFront' : [0.0, 1.0],           
        'shoulder_poserBack' : [0.0, 1.0],
        'shoulder_poserUp' : [0.0, 1.0],
        'shoulder_poserDown' : [0.0, 1.0],
        'clav_40' : [0.0, 8.0],
        'wrist_poserfront' : [0.0, 1.0],
        'wrist_poserback' : [0.0, 1.0],           
        'wrist_poserUp' : [0.0, 1.0],
        'wrist_poserDown' : [0.0, 1.0],
        'index_dn_90' : [0.0, -90.0],
        'mid_dn_90' : [0.0, -90.0],
        'ring_dn_90' : [0.0, -90.0],
        'pinky_dn_90' : [0.0, -90.0],
        'hip_fwd_90' : [0.0, 3.0],
        'knee_140' : [0.0, 90.0, 140.0],
        }

    if pose_name in pose_list:
        sdk_val = sdk_dict[pose_name]
        driver_attr = driver_attr_dict[pose_name]
    else:
        print '%s is a non-standard pose. Using default values.' % pose_name
        sdk_val = [0, -90]
        driver_attr = 'rz'
    
    return driver_attr, sdk_val
    

def create_hookup_network(pose_name, poser_name_list=[], side='lf', angle=None, driver_attr=None, sdk_val=[], blend_node=None, blend_channel=None):
    '''
    Create the network needed to hookup to the pose corrective shape.
    Examples of pose names:
        shoulder_poserUp
        wrist_poserUp    
    '''
    
    print '\n----------------------------------------------'
    print 'Create Hookup Network:\n'
    print 'pose_name = ', pose_name
    print 'side = ', side
    print 'angle = ', angle
    print 'blend_node = ', blend_node
    print 'blend_channel = ', blend_channel    
    print ''  
    
    # Get the standard poser names used for this pose.
    poser_name_standard_list = get_posers_for_pose(pose_name)
    
    # if the user hasn't specified the names of the poser nodes, use the standard ones.
    if not poser_name_list:
        if poser_name_standard_list:
            poser_name_list = list(poser_name_standard_list)
        else:
            maya.mel.eval('warning("No poser nodes were found. Unable to continue.");')
            return        
           
    # Get the standard set driven key values for this pose       
    driver_attr_standard, sdk_val_standard = get_pose_data(pose_name)
 
    # If the user hasn't specified the driver attr and SDK values, use the standard values
    if not driver_attr:
        driver_attr = driver_attr_standard
    
    if not sdk_val:
        sdk_val = sdk_val_standard
    
    # Determine the angle to be used in the setup
    # This is important for shoulder and wrist setups
    if not angle:
        if 'shoulder' in pose_name or 'wrist' in pose_name:
            angle = 90
        else:
            if sdk_val:
                angle = abs(sdk_val[-1])
            else:
                angle = 90
    
    print '\nPoser Names:\n'
    for item in poser_name_list:
        print poser_name_list
    num_poser = len(poser_name_list)

    print ''
    print 'driver_attr = ', driver_attr
    print 'sdk_val = ', sdk_val
    print 'angle = ', angle
    print ''
    
    if 'shoulder' in pose_name:
        create_shoulder_hookup(pose_name, side, angle, driver_attr, sdk_val, poser_name_list, blend_node, blend_channel)    
    elif 'wrist' in pose_name:
        create_wrist_hookup(pose_name, side, angle, driver_attr, sdk_val, poser_name_list, blend_node, blend_channel)
    else:
        create_simple_hookup(pose_name, side, angle, driver_attr, sdk_val, poser_name_list, blend_node, blend_channel)
    
    try:
        cmds.select(blend_node, r=True)
    except Exception, e:
        print e
    
    print '\n'
    print 'Done!'
    
    return    

def create_simple_hookup(pose_name, side='lf', angle=90, driver_attr='rz', sdk_val=[], poser_name_list=[], blend_node=None, blend_channel=None):
    '''
    Create the network needed to hookup the pose corrective shape for simple joints like elbows and knees.
    '''

    print ('\n')
    print ('Create Simple Hookup:\n')
    
    poser_name = ''
    
    if poser_name_list:
        poser_name = poser_name_list[0]
    else:
        maya.mel.eval('warning("No poser nodes were specified. Unable to continue.");')
        return
    
    poser = '%s_%s' % (side, poser_name) 
    print 'poser = ', poser
    
    # set driven key
    if blend_node and blend_channel:
        connect_to_blendshape(poser, driver_attr, blend_node, blend_channel, sdk_val)         
    else:    
        maya.mel.eval('warning("No blendShape specified. Unable to create Set Driven Key.");') 
    
    return    
    
    



def create_shoulder_hookup(pose_name, side='lf', angle=90, driver_attr='rz', sdk_val=[], poser_name_list=[], blend_node=None, blend_channel=None):
    '''
    Create the network needed to hookup the pose corrective shape for an arm pose.
    '''

    print ('\n')
    print ('Create Arm Hookup:\n')
    
    if not angle:
        angle = 90
    print 'angle = ', angle
    
    num_poser = len(poser_name_list)
    poser_list = []
    matrix_list = []
    error_count = 0
    for i in range(num_poser):
        poser_name = poser_name_list[i]
        poser = '%s_%s' % (side, poser_name)
        print '%s.' % i
        print 'poser = ', poser
        if cmds.objExists(poser):
            poser_list.append(poser)    
        else:
            msg = '%s --> NOT FOUND.' % poser
            maya.mel.eval('warning("%s.");' % msg)
            error_count += 1
            continue        
        
        if poser_name == 'hand_loc1':
            poser_description = 'hand_poser'
        else:
            poser_description = poser_name
        print 'poser_description = ', poser_description
        
        matrix = None
        matrix_name = '%s_%s_matrix' % (side, poser_description)
        if cmds.objExists(matrix_name):
            matrix = matrix_name 
        else:
            matrix = cmds.createNode('decomposeMatrix', n='%s_%s_matrix' % (side, poser_description))               
            if cmds.objExists(poser):
                cmds.connectAttr('%s.worldMatrix[0]' % poser, '%s.inputMatrix' % matrix)
            else:
                maya.mel.eval('warning("Poser Node NOT FOUND.");')
        print 'matrix = ', matrix
        
        matrix_list.append(matrix) 

    if error_count > 0:
        msg = '%s poser nodes missing from the file. Unable to continue.' % error_count
        maya.mel.eval('warning("%s.");' % msg)
        return 
    
    # plus minus average nodes
    plus_shoulder = None
    plus_shoulder_name = '%s_shoulder_poser_vector' % side
    if cmds.objExists(plus_shoulder_name):
        plus_shoulder = plus_shoulder_name
    else:
        plus_shoulder = cmds.createNode('plusMinusAverage', n=plus_shoulder_name)
        cmds.setAttr('%s.operation' % plus_shoulder, 2)        
        cmds.connectAttr('%s.outputTranslate' % matrix_list[0], '%s.input3D[0]' % plus_shoulder, f=True)    
        cmds.connectAttr('%s.outputTranslate' % matrix_list[1], '%s.input3D[1]' % plus_shoulder, f=True)         
    print plus_shoulder
    
    plus_pose = None
    plus_pose_name = '%s_%s_vector' % (side, poser_list[-1])
    if cmds.objExists(plus_pose_name):
        plus_pose = plus_pose_name
    else:
        plus_pose = cmds.createNode('plusMinusAverage', n=plus_pose_name) 
        cmds.setAttr('%s.operation' % plus_pose, 2)
        cmds.connectAttr('%s.outputTranslate' % matrix_list[0], '%s.input3D[0]' % plus_pose, f=True)    
        cmds.connectAttr('%s.outputTranslate' % matrix_list[-1], '%s.input3D[1]' % plus_pose, f=True)
    
    print plus_pose
    
    # angle_between
    angle_node = cmds.createNode('angleBetween', n='%s_%s_angleBetween' % (side, pose_name))
    cmds.connectAttr('%s.output3D' % plus_shoulder, '%s.vector2' % angle_node, f=True)
    cmds.connectAttr('%s.output3D' % plus_pose, '%s.vector1' % angle_node, f=True)    
    print angle_node
    
    # multiplyDivide
    multi = cmds.createNode('multiplyDivide', n='%s_%s_multiplyDivide' % (side, pose_name))
    cmds.setAttr('%s.operation' % multi, 2)
    cmds.setAttr('%s.input2X' % multi, angle)
    cmds.connectAttr('%s.angle' % angle_node, '%s.input1X' % multi, f=True)      
    print multi, angle   

    # another minus average node
    plus_extra_name = '%s_%s_pose' % (side, pose_name)
    plus_extra = cmds.createNode('plusMinusAverage', n=plus_extra_name)
    cmds.setAttr('%s.input1D[0]' % plus_extra, k=True, cb=True)
    cmds.setAttr('%s.input1D[1]' % plus_extra, k=True, cb=True)
    cmds.setAttr('%s.input1D[0]' % plus_extra, 1)
    cmds.setAttr('%s.operation' % plus_extra, 2)
    cmds.connectAttr('%s.outputX' % multi, '%s.input1D[1]' % plus_extra, f=True)          
    print plus_extra    
    
    # set driven key
    if blend_node and blend_channel:
        connect_to_blendshape(plus_extra, 'output1D', blend_node, blend_channel, sdk_val)         
    else:    
        maya.mel.eval('warning("No blendShape specified. Unable to create Set Driven Key.");') 
    
    return



def create_wrist_hookup(pose_name, side='lf', angle=90, driver_attr='rz', sdk_val=[], poser_name_list=[], blend_node=None, blend_channel=None):
    '''
    Create the network needed to hookup the pose corrective shape for an arm pose.
    '''

    print ('\n')
    print ('Create Wrist Hookup:\n')
    
    if not angle:
        angle = 90
    print 'angle = ', angle
    
    num_poser = len(poser_name_list)
    poser_list = []
    matrix_list = []
    error_count = 0
    for i in range(num_poser):
        poser_name = poser_name_list[i]
        poser = '%s_%s' % (side, poser_name)
        print '%s.' % i
        print 'poser = ', poser
        if cmds.objExists(poser):
            poser_list.append(poser)    
        else:
            msg = '%s --> NOT FOUND.' % poser
            maya.mel.eval('warning("%s.");' % msg)
            error_count += 1
            continue
        
        if poser_name == 'hand_loc1':
            poser_description = 'hand_poser'
        else:
            poser_description = poser_name
        print 'poser_description = ', poser_description
        
        matrix = None
        matrix_name = '%s_%s_matrix' % (side, poser_description)
        if cmds.objExists(matrix_name):
            matrix = matrix_name 
        else:
            matrix = cmds.createNode('decomposeMatrix', n='%s_%s_matrix' % (side, poser_description))               
            if cmds.objExists(poser):
                cmds.connectAttr('%s.worldMatrix[0]' % poser, '%s.inputMatrix' % matrix)
            else:
                maya.mel.eval('warning("Poser Node NOT FOUND.");')
        print 'matrix = ', matrix
        
        matrix_list.append(matrix) 

    if error_count > 0:
        msg = '%s poser nodes missing from the file. Unable to continue.' % error_count
        maya.mel.eval('warning("%s.");' % msg)
        return
    
    # plus minus average nodes
    plus_wrist = None
    plus_wrist_name = '%s_wrist_poser_vector' % side
    if cmds.objExists(plus_wrist_name):
        plus_wrist = plus_wrist_name
    else:
        plus_wrist = cmds.createNode('plusMinusAverage', n=plus_wrist_name)
        cmds.setAttr('%s.operation' % plus_wrist, 2)        
        cmds.connectAttr('%s.outputTranslate' % matrix_list[0], '%s.input3D[1]' % plus_wrist, f=True)    
        cmds.connectAttr('%s.outputTranslate' % matrix_list[1], '%s.input3D[0]' % plus_wrist, f=True)         
    print plus_wrist
    
    plus_pose = None
    plus_pose_name = '%s_%s_vector' % (side, poser_list[-1])
    if cmds.objExists(plus_pose_name):
        plus_pose = plus_pose_name
    else:
        plus_pose = cmds.createNode('plusMinusAverage', n=plus_pose_name) 
        cmds.setAttr('%s.operation' % plus_pose, 2)
        cmds.connectAttr('%s.outputTranslate' % matrix_list[0], '%s.input3D[1]' % plus_pose, f=True)    
        cmds.connectAttr('%s.outputTranslate' % matrix_list[-1], '%s.input3D[0]' % plus_pose, f=True) 
    print plus_pose
    
    # angle_between
    angle_node = cmds.createNode('angleBetween', n='%s_%s_angleBetween' % (side, pose_name))
    cmds.connectAttr('%s.output3D' % plus_wrist, '%s.vector1' % angle_node, f=True)
    cmds.connectAttr('%s.output3D' % plus_pose, '%s.vector2' % angle_node, f=True)    
    print angle_node
    
    # multiplyDivide
    multi = cmds.createNode('multiplyDivide', n='%s_%s_multiplyDivide' % (side, pose_name))
    cmds.setAttr('%s.operation' % multi, 2)
    cmds.setAttr('%s.input2X' % multi, angle)
    cmds.connectAttr('%s.angle' % angle_node, '%s.input1X' % multi, f=True)      
    print multi, angle   

    # another minus average node
    plus_extra_name = '%s_%s_pose' % (side, pose_name)
    plus_extra = cmds.createNode('plusMinusAverage', n=plus_extra_name)
    cmds.setAttr('%s.input1D[0]' % plus_extra, k=True, cb=True)
    cmds.setAttr('%s.input1D[1]' % plus_extra, k=True, cb=True)
    cmds.setAttr('%s.input1D[0]' % plus_extra, 1)
    cmds.setAttr('%s.operation' % plus_extra, 2)
    cmds.connectAttr('%s.outputX' % multi, '%s.input1D[1]' % plus_extra, f=True)          
    print plus_extra    
    
    # set driven key
    if blend_node and blend_channel:
        connect_to_blendshape(plus_extra, 'output1D', blend_node, blend_channel, sdk_val)         
    else:    
        maya.mel.eval('warning("No blendShape specified. Unable to create Set Driven Key.");') 
    
    return


def connect_to_blendshape(src_item, src_attr, blend_node, blend_channel, sdk_val):
    '''
    Connect the network to the blendShape node.
    IMPORTANT: The set driven keys will have to be adjusted by hand after the connection is made.
    '''
    
    print '\nCreating Set Driven Keys...\n'
    #print src_item, src_attr, blend_node, blend_channel, sdk_val
    
    num_val = 0
    if sdk_val:
        num_val = len(sdk_val)
    
    val_list = []
    if num_val == 2:
        val_list = [0, 1]       
    elif num_val == 3:
        val_list = [0, sdk_val[1]/sdk_val[2], 1]
    else:
        val_list = list(sdk_val)

    print 'num_val = ', num_val
    print 'sdk_val = ', sdk_val
    print 'val_list = ', val_list
        
    error_count = 0
    if not cmds.objExists('%s.%s' % (src_item, src_attr)):
        print '%s.%s --> NOT FOUND.' % (src_item, src_attr)
        error_count += 1
    if not cmds.objExists('%s.%s' % (blend_node, blend_channel)):
        print '%s.%s --> NOT FOUND.' % (blend_node, blend_channel)
        error_count += 1
    
    # Create the set driven keys
    for i in range(num_val):
        driver_val = sdk_val[i]
        val = val_list[i] 
        cmds.setDrivenKeyframe(blend_node, at=blend_channel, cd='%s.%s' % (src_item, src_attr), dv=driver_val, v=val, itt='linear', ott='linear')
    
    if error_count > 0:
        maya.mel.eval('warning("Errors may have occurred while creating the set driven keys.");')           
    else:
        print 'Set driven keys created.'
    
    return



    
    
